{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Welcome to the 30-minute Flower federated learning tutorial!\n",
    "\n",
    "In this tutorial you will implement your first Federated Learning project using [Flower](https://flower.ai/).\n",
    "\n",
    "> You can find a similar tutorial, but using the preferred `flwr run` CLI command to launch your experiments in [the Flower Documentation](https://flower.ai/docs/framework/tutorial-quickstart-pytorch.html). Running Flower from the CLI is recommended. Not all Flower features are available when skipping `flwr run` and when using the Pythonic entrypoint of the Simulation Runtime (i.e. calling directly [run_simulation](https://flower.ai/docs/framework/ref-api/flwr.simulation.run_simulation.html)).\n",
    "\n",
    "🧑‍🏫 This tutorial starts at zero and expects no familiarity with federated learning. Only a basic understanding of data science and Python programming is assumed. A minimal understanding of ML is not required but if you already know about it, nothing is stopping your from modifying this code as you see fit!\n",
    "\n",
    "> [Star Flower on GitHub](https://github.com/adap/flower) ⭐️ and join the Flower community on Flower Discuss and the Flower Slack to connect, ask questions, and get help:\n",
    "> - [Join Flower Discuss](https://discuss.flower.ai/) We'd love to hear from you in the `Introduction` topic! If anything is unclear, post in `Flower Help - Beginners`.\n",
    "> - [Join Flower Slack](https://flower.ai/join-slack) We'd love to hear from you in the `#introductions` channel! If anything is unclear, head over to the `#questions` channel.\n",
    "\n",
    "Let's get started!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Environment Setup\n",
    "\n",
    "Now let's really begin with this tutorial!\n",
    "\n",
    "To start working with Flower, very little is required once you have activated your Python environment (e.g. via `conda`, `virtualenv`, `pyenv`, etc). If you are running this code on Colab, there is really nothing to do except to install Flower and other dependencies. The steps below have been verified to run in Colab. Let's first, install Flower, then the ML framework of your choice and extra dependencies you might want to use.\n",
    "\n",
    "## Installing Flower\n",
    "\n",
    "You can install flower very conveniently from `pip`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q \"flwr[simulation]\" flwr-datasets"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be using the _simulation_ runtime in Flower, which allows you to run a large number of clients without the overheads of manually managing devices. This is achieved via the `Simulation Engine`, the core component in Flower to run simulations efficiently."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install your ML framework\n",
    "\n",
    "Flower is agnostic to your choice of ML Framework. Flower works with `PyTorch`, `Tensorflow`, `NumPy`, `🤗 Transformers`, `MLX`, `JAX`, `scikit-learn`, `fastai`, `Pandas`. Flower also supports all major platforms: `iOS`, `Android` and plain `C++`. You can find a _quickstart- example for each of the above in the [Flower Repository](https://github.com/adap/flower/tree/main/examples) inside the `examples/` directory. And check the [Flower Documentation](https://flower.ai/docs/) for even more learning materials.\n",
    "\n",
    "In this tutorial we are going to use PyTorch, uncomment the line below if you haven't installed PyTorch in your system. In this tutorial we'll use a small model so using CPU only training will suffice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "7192138a-8c87-4d9a-f726-af1038ad264c"
   },
   "outputs": [],
   "source": [
    "# you might see a warning after running the command below, this can be ignored\n",
    "# if you are running this outside Colab, you probably need to adjust the command below\n",
    "# %pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cpu"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to install some other dependencies you are likely familiar with. We'll use these to make plots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "58b7af77-609f-4118-bd5b-5629a4b5a296"
   },
   "outputs": [],
   "source": [
    "%pip install matplotlib"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Centralised training: the old way of doing ML"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's begin by creating a simple (but complete) training loop as it is commonly done in centralised setups. Starting our tutorial in this way will allow us to very clearly identify which parts of a typical ML pipeline are common to both centralised and federated training and which ones are poles a part.\n",
    "\n",
    "For this tutorial we'll design a image classification pipeline for [MNIST digits](https://en.wikipedia.org/wiki/MNIST_database) and using a simple CNN model as the network to train. The MNIST dataset is comprised of `28x28` greyscale images with digits from 0 to 9 (i.e. 10 classes in total)\n",
    "\n",
    "\n",
    "## A dataset\n",
    "\n",
    "Let's begin by constructing the dataset. We will use 🤗HuggingFace Datasets to download MNIST. We will prepare a function that will be use later to apply standard normalization transformations from `TorchVision` and create the dataloaders for the `train` and `test` partitions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we naturally first need to import torch and torchvision\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.transforms import ToTensor, Normalize, Compose\n",
    "from datasets import load_dataset\n",
    "\n",
    "\n",
    "def get_mnist_dataloaders(mnist_dataset, batch_size: int):\n",
    "    pytorch_transforms = Compose([ToTensor(), Normalize((0.1307,), (0.3081,))])\n",
    "\n",
    "    # Prepare transformation functions\n",
    "    def apply_transforms(batch):\n",
    "        batch[\"image\"] = [pytorch_transforms(img) for img in batch[\"image\"]]\n",
    "        return batch\n",
    "\n",
    "    mnist_train = mnist_dataset[\"train\"].with_transform(apply_transforms)\n",
    "    mnist_test = mnist_dataset[\"test\"].with_transform(apply_transforms)\n",
    "\n",
    "    # Construct PyTorch dataloaders\n",
    "    trainloader = DataLoader(mnist_train, batch_size=batch_size, shuffle=True)\n",
    "    testloader = DataLoader(mnist_test, batch_size=batch_size)\n",
    "    return trainloader, testloader\n",
    "\n",
    "\n",
    "# Download dataset\n",
    "mnist = load_dataset(\"ylecun/mnist\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can have a quick overview of our datasets by just typing the object on the command line. For instance, below you can see the sizes of both the `train` and `test` partitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "f10b649f-3cee-4e86-c7ff-94bd1fd3e082"
   },
   "outputs": [],
   "source": [
    "mnist"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a more insightful visualisation. First let's see the distribution over the labels by constructing a histogram. Then, let's visualise some training examples !"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 490
    },
    "outputId": "c8d0f4c0-60cd-4c58-bc91-3b061dae8046"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from collections import Counter\n",
    "\n",
    "\n",
    "# construct histogram\n",
    "all_labels = mnist[\"train\"][\"label\"]\n",
    "all_label_counts = Counter(all_labels)\n",
    "\n",
    "# visualise histogram\n",
    "bar = plt.bar(all_label_counts.keys(), all_label_counts.values())\n",
    "_ = plt.bar_label(bar)\n",
    "\n",
    "# plot formatting\n",
    "_ = plt.xticks([label for label in all_label_counts.keys()])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualise 32 images from the dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import io\n",
    "\n",
    "\n",
    "def visualise_n_random_examples(trainset_, n: int, verbose: bool = True):\n",
    "    trainset_data = [\n",
    "        Image.open(io.BytesIO(entry[0].as_py())) for entry in trainset_.data[0]\n",
    "    ]\n",
    "    idx = list(range(len(trainset_data)))\n",
    "    random.shuffle(idx)\n",
    "    idx = idx[:n]\n",
    "    if verbose:\n",
    "        print(f\"will display images with idx: {idx}\")\n",
    "\n",
    "    # construct canvas\n",
    "    num_cols = 8\n",
    "    num_rows = int(np.ceil(len(idx) / num_cols))\n",
    "    fig, axs = plt.subplots(figsize=(16, num_rows * 2), nrows=num_rows, ncols=num_cols)\n",
    "\n",
    "    # display images on canvas\n",
    "    for c_i, i in enumerate(idx):\n",
    "        axs.flat[c_i].imshow(trainset_data[i], cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 715
    },
    "outputId": "4e0988a8-388d-4acf-882b-089e4ea887bf"
   },
   "outputs": [],
   "source": [
    "# it is likely that the plot this function will generate looks familiar to other plots you might have generated before\n",
    "# or you might have encountered in other tutorials. So far, we aren't doing anything new, Federated Learning will start soon!\n",
    "visualise_n_random_examples(mnist[\"train\"], n=32)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A CNN architecture\n",
    "\n",
    "This tutorial is not so much about novel architectural designs so we keep things simple and make use of a typical CNN that is adequate for the MNIST image classification task.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self, num_classes: int) -> None:\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 6, 5)\n",
    "        self.pool = nn.MaxPool2d(2, 2)\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
    "        self.fc1 = nn.Linear(16 * 4 * 4, 120)\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.fc3 = nn.Linear(84, num_classes)\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x = self.pool(F.relu(self.conv2(x)))\n",
    "        x = x.view(-1, 16 * 4 * 4)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly to what we did with the dataset you could inspect the model in various ways. We can, for instance, count the number of model parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "67d01ab4-cdd9-4661-8f01-eaa9aabf786d"
   },
   "outputs": [],
   "source": [
    "model = Net(num_classes=10)\n",
    "num_parameters = sum(value.numel() for value in model.state_dict().values())\n",
    "print(f\"{num_parameters = }\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Training Loop\n",
    "\n",
    "A minimal training loop in PyTorch can be constructed with three functions:\n",
    "*  `train()` that will train the model given a dataloader.\n",
    "* `test()` that will be used to evaluate the performance of the model on held-out data, e.g., a training set.\n",
    "* `run_centralised()` which will define additional elements (e.g. the optimiser) and run the training loop over N epochs.\n",
    "\n",
    "Let's construct these functions!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_fn(net, trainloader, optimizer, device=\"cpu\"):\n",
    "    \"\"\"Train the network on the training set.\"\"\"\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    net.to(device)\n",
    "    net.train()\n",
    "    loss_accum = 0.0\n",
    "    for batch in trainloader:\n",
    "        images, labels = batch[\"image\"].to(device), batch[\"label\"].to(device)\n",
    "        optimizer.zero_grad()\n",
    "        loss = criterion(net(images), labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        loss_accum += loss.item()\n",
    "    return loss_accum / len(trainloader)\n",
    "\n",
    "\n",
    "def test_fn(net, testloader, device):\n",
    "    \"\"\"Validate the network on the entire test set.\"\"\"\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    correct, loss = 0, 0.0\n",
    "    net.to(device)\n",
    "    net.eval()\n",
    "    with torch.no_grad():\n",
    "        for batch in testloader:\n",
    "            images, labels = batch[\"image\"].to(device), batch[\"label\"].to(device)\n",
    "            outputs = net(images)\n",
    "            loss += criterion(outputs, labels).item()\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    accuracy = correct / len(testloader.dataset)\n",
    "    return loss, accuracy\n",
    "\n",
    "\n",
    "def run_centralised(\n",
    "    trainloader, testloader, epochs: int, lr: float, momentum: float = 0.9\n",
    "):\n",
    "    \"\"\"A minimal (but complete) training loop\"\"\"\n",
    "\n",
    "    # instantiate the model\n",
    "    model = Net(num_classes=10)\n",
    "\n",
    "    # Discover device\n",
    "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "    model.to(device)\n",
    "\n",
    "    # define optimiser with hyperparameters supplied\n",
    "    optim = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)\n",
    "\n",
    "    # train for the specified number of epochs\n",
    "    for e in range(epochs):\n",
    "        print(f\"Training epoch {e} ...\")\n",
    "        train_fn(model, trainloader, optim, device)\n",
    "\n",
    "    # training is completed, then evaluate model on the test set\n",
    "    loss, accuracy = test_fn(model, testloader, device)\n",
    "    print(f\"{loss = }\")\n",
    "    print(f\"{accuracy = }\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's run this for 5 epochs (you'll see it reaching close to 99% accuracy -- as expected from a centralised setup with the MNIST dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "e8d9b429-178d-4924-e82f-4d4e52863788"
   },
   "outputs": [],
   "source": [
    "# Construct dataloaders\n",
    "trainloader, testloader = get_mnist_dataloaders(mnist, batch_size=32)\n",
    "\n",
    "# Run the centralised training\n",
    "run_centralised(trainloader, testloader, epochs=3, lr=0.01)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above centralised formulation has worked just fine for some applications and to showcase the potential of AI in a variety of scenarios. However, as was discussed earlier in the session, centralised training is unsuitable for a larger range of settings where information cannot be collected in order to build a unified (centralised) dataset.\n",
    "\n",
    "# The Future of AI is Federated\n",
    "\n",
    "What are the key differences between Federated Learning and Centralised Training? If you could only pick you, probably you'd say:\n",
    "* Federated Learning is distributed -- the model is trained on-device by the participating clients.\n",
    "* Data remains private and is owned by a specific _client_ -- the data is never sent to the central server.\n",
    "\n",
    "There are several more differences. But the above two are the main ones to always consider and that are common to all flavours of Federated Learning (e.g. _cross-device_ or _cross-silo_). The rest of this tutorial is going to focus on transforming the code we have written so far for the centralised setting and construct a Federated Learning pipeline using Flower and PyTorch.\n",
    "\n",
    "Let's begin! 🚀"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## One Client, One Data Partition\n",
    "\n",
    "To start designing a Federated Learning pipeline, we need to meet one of the key properties in FL: each client has its own data partition. To accomplish this with the MNIST dataset, we are going to generate N random partitions, where N is the total number of clients in our FL system, using [Flower Datasets](https://flower.ai/docs/datasets/). Let's create 100 partitions with the [IidPartitioner](https://flower.ai/docs/datasets/ref-api/flwr_datasets.partitioner.IidPartitioner.html#flwr_datasets.partitioner.IidPartitioner) -- note there are many more [partitioners](https://flower.ai/docs/datasets/ref-api/flwr_datasets.partitioner.html) to choose from."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flwr_datasets import FederatedDataset\n",
    "from flwr_datasets.partitioner import IidPartitioner\n",
    "\n",
    "NUM_PARTITIONS = 100\n",
    "\n",
    "partitioner = IidPartitioner(num_partitions=NUM_PARTITIONS)\n",
    "# Let's partition the \"train\" split of the MNIST dataset\n",
    "# The MNIST dataset will be downloaded if it hasn't been already\n",
    "fds = FederatedDataset(dataset=\"ylecun/mnist\", partitioners={\"train\": partitioner})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Accessing individual partitions can be done like this. The return object can then be passed to a dataloader for training or evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We could load a single partition like this\n",
    "partition_0 = fds.load_partition(0)\n",
    "partition_0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Flower Datasets` comes with built-in visualization tools that help you gain insights into how the dataset (in this case MNIST) has been partitioned. Let's create a bar plot to visualize the number of labels of each class that every client's partition contains. Note we are only visualising the first 30 clients purely so the plot remains readable.\n",
    "\n",
    "> There are many more types of plots you can generate with Flower Datasets. Check the [Visualization tutorial](https://flower.ai/docs/datasets/tutorial-visualize-label-distribution.html). Feel free to try other partitioning schemes and you'll see how the visualization changes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flwr_datasets.visualization import plot_label_distributions\n",
    "\n",
    "fig, ax, df = plot_label_distributions(\n",
    "    partitioner,\n",
    "    label_name=\"label\",\n",
    "    plot_type=\"bar\",\n",
    "    size_unit=\"absolute\",\n",
    "    partition_id_axis=\"x\",\n",
    "    legend=True,\n",
    "    verbose_labels=True,\n",
    "    max_num_partitions=30,  # Note we are only showing the first 30 so the plot remains readable\n",
    "    title=\"Per Partition Labels Distribution\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the histogram of this partition is a bit different from the one we obtained at the beginning, where we took the entire dataset into consideration. Because our data partitions are artificially constructed by sampling the MNIST dataset in an IID fashion, our Federated Learning example will not face severe _data heterogeneity_ issues (which is a fairly [active research topic](https://arxiv.org/abs/1912.04977))."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constructing Messages\n",
    "\n",
    "In Flower, the server and clients communicate by sending and receiving [Message](https://flower.ai/docs/framework/ref-api/flwr.app.Message.html#flwr.app.Message)\n",
    "objects. A `Message` carries a [RecordDict](https://flower.ai/docs/framework/ref-api/flwr.app.RecordDict.html#flwr.app.RecordDict) as its main payload. The `RecordDict`\n",
    "is like a Python dictionary that can contain multiple records of different types. There\n",
    "are three main types of records:\n",
    "\n",
    "- [ArrayRecord](https://flower.ai/docs/framework/ref-api/flwr.app.ArrayRecord.html): Contains model parameters as a dictionary of NumPy arrays\n",
    "- [MetricRecord](https://flower.ai/docs/framework/ref-api/flwr.app.MetricRecord.html): Contains training or evaluation metrics as a dictionary of\n",
    "  integers, floats, lists of integers, or lists of floats.\n",
    "- [ConfigRecord](https://flower.ai/docs/framework/ref-api/flwr.app.ConfigRecord.html#flwr.app.ConfigRecord): Contains configuration parameters as a dictionary of integers,\n",
    "  floats, strings, booleans, or bytes. Lists of these types are also supported.\n",
    "\n",
    "Let's see a few examples of how to work with these types of records and, ultimately,\n",
    "construct a `RecordDict` that can be sent over a `Message`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flwr.app import ArrayRecord, MetricRecord, ConfigRecord, RecordDict\n",
    "\n",
    "# ConfigRecord can be used to communicate configs between ServerApp and ClientApp\n",
    "# They can hold scalars, but also strings and booleans\n",
    "config = ConfigRecord(\n",
    "    {\"batch_size\": 32, \"use_augmentation\": True, \"data-path\": \"/my/dataset\"}\n",
    ")\n",
    "\n",
    "# MetricRecord objects expect scalar-based metrics (i.e. int/float/list[int]/list[float])\n",
    "# By limiting the types Flower can aggregate MetricRecord objects automatically\n",
    "metrics = MetricRecord({\"accuracy\": 0.9, \"losses\": [0.1, 0.001], \"perplexity\": 2.31})\n",
    "\n",
    "# ArrayRecord objects are designed to communicate arrays/tensors/weights from ML models\n",
    "model = Net(num_classes=10)\n",
    "array_record = ArrayRecord(model.state_dict())  # for a PyTorch model\n",
    "# Then you can convert the ArrayRecord back into a PyTorch state_dict with:\n",
    "state_dict = array_record.to_torch_state_dict()\n",
    "\n",
    "# A RecordDict is like a dictionary that holds named records.\n",
    "# This is the main payload of a Message\n",
    "rd = RecordDict({\"my-config\": config, \"metrics\": metrics, \"my-model\": array_record})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Check the [ArrayRecord](https://flower.ai/docs/framework/ref-api/flwr.app.ArrayRecord.html#flwr.app.ArrayRecord) documentation to learn how to use models from other ML frameworks other than PyTorch."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's define how our FL clients will behave\n",
    "\n",
    "## Defining a Flower `ClientApp`\n",
    "\n",
    "\n",
    "You can think of a client in FL as an entity that owns some data and trains a model using this data. The caveat is that the model is being trained _collaboratively_ in a federation by multiple clients (sometimes numbering in the hundreds of thousands) and, in most instances of FL, is sent by a central server running in a `ServerApp` (more on this later).\n",
    "\n",
    "We can define how the `ClientApp` performs training by wrapping a function with the\n",
    "`@app.train()` decorator. In this case we name this function `train` because we'll\n",
    "use it to train the model on the local data. The function always expects two arguments:\n",
    "\n",
    "- A [Message](https://flower.ai/docs/framework/ref-api/flwr.app.Message.html#flwr.app.Message): The message received from the server. It contains the model\n",
    "  parameters and any other configuration information sent by the server.\n",
    "- A [Context](https://flower.ai/docs/framework/ref-api/flwr.app.Context.html#flwr.app.Context): The context object that contains information about the node\n",
    "  executing the `ClientApp` and about the current run. \n",
    "\n",
    "> **Note**: When running Flower on a notebook, the `run_config` in the `Context` is always empty. It is recommended to run Flower on the terminal via `flwr run` instead of in a notebook/Colab. Check the [Flower Tutorial](https://flower.ai/docs/framework/tutorial-series-get-started-with-flower-pytorch.html) to learn more.\n",
    "\n",
    "In Flower, `ClientApps` are\n",
    "ephemeral objects that get instantiated for the execution of one `Message` and\n",
    "destroyed when a reply is communicated back to the server.\n",
    "\n",
    "Let's see an implementation of `ClientApp` that uses the previously defined PyTorch\n",
    "CNN model, applies the parameters received from the `ServerApp` via the message, loads\n",
    "its local data, trains the model with it (using the `train_fn` function), and\n",
    "generates a reply `Message` containing the updated model parameters as well as some\n",
    "metrics of interest."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's define our Flower Client class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flwr.app import Context, Message\n",
    "from flwr.clientapp import ClientApp\n",
    "\n",
    "# Flower ClientApp\n",
    "client_app = ClientApp()\n",
    "\n",
    "\n",
    "@client_app.train()\n",
    "def train(msg: Message, context: Context):\n",
    "    \"\"\"Train the model on local data.\"\"\"\n",
    "\n",
    "    # Load the model and initialize it with the received weights\n",
    "    model = Net(num_classes=10)\n",
    "    model.load_state_dict(msg.content[\"arrays\"].to_torch_state_dict())\n",
    "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "    model.to(device)\n",
    "\n",
    "    # Load the data\n",
    "    partition_id = context.node_config[\"partition-id\"]\n",
    "    partition = fds.load_partition(partition_id, \"train\")\n",
    "    # partition into train/validation\n",
    "    partition_train_val = partition.train_test_split(test_size=0.1, seed=42)\n",
    "\n",
    "    # Let's use the function defined earlier to construct the dataloaders\n",
    "    # and apply the dataset transformations\n",
    "    trainloader, _ = get_mnist_dataloaders(partition_train_val, batch_size=32)\n",
    "\n",
    "    # Let's defined the optimizer\n",
    "    lr = msg.content[\"config\"][\"lr\"]\n",
    "    optim = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)\n",
    "\n",
    "    # Call the training function\n",
    "    train_loss = train_fn(\n",
    "        model,\n",
    "        trainloader,\n",
    "        optim,\n",
    "        device,\n",
    "    )\n",
    "\n",
    "    # Construct and return reply Message\n",
    "    model_record = ArrayRecord(model.state_dict())\n",
    "    metrics = {\n",
    "        \"train_loss\": train_loss,\n",
    "        \"num-examples\": len(trainloader.dataset),\n",
    "    }\n",
    "    metric_record = MetricRecord(metrics)\n",
    "    content = RecordDict({\"arrays\": model_record, \"metrics\": metric_record})\n",
    "    return Message(content=content, reply_to=msg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In a typical federated learning setup, the `ClientApp` would also implement an\n",
    "`@app.evaluate()` function to evaluate the model received from the `ServerApp` on\n",
    "local validation data. This is especially useful to monitor the performance of the\n",
    "global model on each client during training. The implementation of the `evaluate`\n",
    "function is very similar to the `train` function, except that it calls the `test_fn`\n",
    "function defined earlier in this tutorial (which implements the PyTorch evaluation loop)\n",
    "and it returns a `Message` containing only a `MetricRecord` with the evaluation\n",
    "metrics (no `ArrayRecord` because the model parameters are not updated during\n",
    "evaluation). Here's how the `evaluate` function looks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@client_app.evaluate()\n",
    "def evaluate(msg: Message, context: Context):\n",
    "    \"\"\"Evaluate the model on local data.\"\"\"\n",
    "\n",
    "    # Load the model and initialize it with the received weights\n",
    "    model = Net(num_classes=10)\n",
    "    model.load_state_dict(msg.content[\"arrays\"].to_torch_state_dict())\n",
    "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "    model.to(device)\n",
    "\n",
    "    # Load the data\n",
    "    partition_id = context.node_config[\"partition-id\"]\n",
    "    partition = fds.load_partition(partition_id, \"train\")\n",
    "    # partition into train/validation\n",
    "    partition_train_val = partition.train_test_split(test_size=0.1, seed=42)\n",
    "\n",
    "    # Let's use the function defined earlier to construct the dataloaders\n",
    "    # and apply the dataset transformations. Let's use the validation set here\n",
    "    _, valloader = get_mnist_dataloaders(partition_train_val, batch_size=32)\n",
    "\n",
    "    # Call the evaluation function\n",
    "    eval_loss, eval_acc = test_fn(\n",
    "        model,\n",
    "        valloader,\n",
    "        device,\n",
    "    )\n",
    "\n",
    "    # Construct and return reply Message\n",
    "    metrics = {\n",
    "        \"eval_loss\": eval_loss,\n",
    "        \"eval_acc\": eval_acc,\n",
    "        \"num-examples\": len(valloader.dataset),\n",
    "    }\n",
    "    metric_record = MetricRecord(metrics)\n",
    "    content = RecordDict({\"metrics\": metric_record})\n",
    "    return Message(content=content, reply_to=msg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the `evaluate` implementation is nearly identical to the `train`\n",
    "implementation, except that it calls the `test_fn` function instead of the\n",
    "`train_fn` function and it returns a `Message` containing only a `MetricRecord`\n",
    "with metrics relevant to evaluation (`eval_loss`, `eval_acc` -- both scalars). We\n",
    "also need to include the `num-examples` key in the metrics so the server can aggregate\n",
    "the evaluation metrics correctly.\n",
    "\n",
    "Now that a `ClientApp` is fully defined, let's create its counterpart: the `ServerApp`."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining a Flower `ServerApp`\n",
    "\n",
    "On the server side, we need to configure a strategy which encapsulates the federated\n",
    "learning approach/algorithm, for example, *Federated Averaging* (FedAvg). Flower has a\n",
    "number of [built-in strategies](https://flower.ai/docs/framework/ref-api/flwr.serverapp.strategy.html#module-flwr.serverapp.strategy), but we can also use our own strategy implementations to\n",
    "customize nearly all aspects of the federated learning approach. For this tutorial, we\n",
    "use the built-in [FedAvg](https://flower.ai/docs/framework/ref-api/flwr.serverapp.strategy.FedAvg.html#flwr.serverapp.strategy.FedAvg) implementation and customize it slightly by specifying the\n",
    "fraction of connected nodes to involve in a round of training.\n",
    "\n",
    "To construct a |serverapp_link|_, we define its `@app.main()` method. This method\n",
    "receives as input arguments:\n",
    "\n",
    "- a `Grid` object that will be used to interface with the nodes running the\n",
    "  `ClientApp` to involve them in a round of train/evaluate/query or other tasks.\n",
    "- a `Context` object that provides access to the run configuration.\n",
    "\n",
    "> Recall that, just like for the `ClientApp` the run config is empty when running Flower without `flwr run`.\n",
    "\n",
    "Before launching the strategy via the [start()](https://flower.ai/docs/framework/ref-api/flwr.serverapp.strategy.Strategy.html#flwr.serverapp.strategy.Strategy.start) method, we want to\n",
    "initialize the global model. This will be the model that gets sent to the `ClientApp`\n",
    "running on the clients in the first round of federated learning. We can do this by\n",
    "creating an instance of the model (`Net`), extracting the parameters from its\n",
    "`state_dict`, and constructing an `ArrayRecord` with them. We can then make it\n",
    "available to the strategy via the `initial_arrays` argument of the `start()` method.\n",
    "\n",
    "We can also optionally pass to the `start()` method a `ConfigRecord` containing\n",
    "settings that we would like to communicate to the clients. These will be sent as part of\n",
    "the `Message` that also carries the model parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flwr.app import Context\n",
    "from flwr.serverapp import Grid, ServerApp\n",
    "from flwr.serverapp.strategy import FedAvg\n",
    "\n",
    "server_app = ServerApp()\n",
    "\n",
    "\n",
    "@server_app.main()\n",
    "def main(grid: Grid, context: Context) -> None:\n",
    "    \"\"\"Main entry point for the ServerApp.\"\"\"\n",
    "\n",
    "    # Load global model\n",
    "    global_model = Net(num_classes=10)\n",
    "    arrays = ArrayRecord(global_model.state_dict())\n",
    "\n",
    "    # Initialize FedAvg strategy\n",
    "    # Sample half of the clients for federated training and\n",
    "    # all of the clients for federated evaluation\n",
    "    strategy = FedAvg(fraction_train=0.5, fraction_evaluate=1.0)\n",
    "\n",
    "    # Start strategy, run FedAvg for `num_rounds`\n",
    "    # We define a `ConfigRecord` with a learning rate of 0.1\n",
    "    # This will be sent to all clients at the beginning of each round of training\n",
    "    result = strategy.start(\n",
    "        grid=grid,\n",
    "        initial_arrays=arrays,\n",
    "        train_config=ConfigRecord({\"lr\": 0.1}),\n",
    "        num_rounds=5,\n",
    "    )\n",
    "\n",
    "    # Save final model to disk\n",
    "    print(\"\\nSaving final model to disk...\")\n",
    "    state_dict = result.arrays.to_torch_state_dict()\n",
    "    torch.save(state_dict, \"final_model.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After the execution of the `start()` method, the `ServerApp` will save the `state_dict` of the resulting model to disk."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Launching the Simulation\n",
    "\n",
    "With both `ClientApp` and `ServerApp` ready, we can launch the simulation. Pass both apps to the `run_simulation()` function and specify the number of `supernodes` (this is a more general term used in Flower to refer to individual \"nodes\" or \"clients\"). We earlier partitioned the dataset into 100 partitions, one for each supernode. So we indicate that `num_supernodes=100`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flwr.simulation import run_simulation\n",
    "\n",
    "run_simulation(\n",
    "    server_app=server_app, client_app=client_app, num_supernodes=NUM_PARTITIONS\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus: Create your own Strategy\n",
    "\n",
    "> Find an extended version of this tutorial in the [how-to use a Flower Strategy](https://flower.ai/docs/framework/tutorial-series-use-a-federated-learning-strategy-pytorch.html) page.\n",
    "\n",
    "Flower strategies can be extended easily to suit your FL setups or your preferred workflows whether you use Flower for research or in production. In this final section, you'll learn how to create a custom strategy that behaves just like `FedAvg` but extends the functionality to facilitate LR schedulling. Then, the `ServerApp` will make a plot once the execution of the `start()` method.\n",
    "\n",
    "Let's call this strategy `FedAvgCustom`. We'll use it to also showcase how to use the `evaluate_fn` callback, a convenient way to do centralised evaluation of the global model after each round. Note this functionality is user-defined since it requires a choice in terms of ML-framework. (Recall that Flower is framework agnostic).\n",
    "\n",
    "> This being said, centralised evaluation of the global model is only possible if there exists a centralised dataset that somewhat follows a similar distribution as the data that's spread across clients. In some cases having such centralised dataset for validation is not possible, so the only solution is to federate the evaluation of the _global model_. This is the default behaviour in Flower. If you don't specify the `evaluate_fn` argument in your strategy, then, centralised global evaluation won't be performed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `FedAvgCustom` defined below only modifies the `configure_train` method of the `FedAvg` strategy. This method is used to prepare the messages that are going to be sent to the `ClientApps`. These messages contain two main payloads: an `ArrayRecord` representing the model to federate; and a `ConfigRecord` containing settings/configs relevant for the `ClientApps` participanting in the given round. Below, we simply apply learning rate decay every 5 rounds by halving the LR found in the config record."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Iterable\n",
    "from flwr.serverapp import Grid\n",
    "from flwr.serverapp.strategy import FedAvg\n",
    "from flwr.app import ArrayRecord, ConfigRecord, Message\n",
    "\n",
    "\n",
    "class FedAvgCustom(FedAvg):\n",
    "    def configure_train(\n",
    "        self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid\n",
    "    ) -> Iterable[Message]:\n",
    "        \"\"\"Configure the next round of federated training and maybe do LR decay.\"\"\"\n",
    "        # Decrease learning rate by a factor of 0.5 every 5 rounds\n",
    "        if server_round % 5 == 0 and server_round > 0:\n",
    "            config[\"lr\"] *= 0.5\n",
    "            print(\"LR decreased to:\", config[\"lr\"])\n",
    "        # Pass the updated config and the rest of arguments to the parent class\n",
    "        return super().configure_train(server_round, arrays, config, grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We’ve seen how federated evaluation works on the client side (i.e., by implementing a function wrapped with the `@app.evaluate` decorator in your `ClientApp`). Now let’s see how we can evaluate the aggregated model parameters on the server side.\n",
    "\n",
    "To do so, we need to create a new function that we can name `central_evaluate`. This function is a callback that will be passed to the `start` method of our strategy. This means that the strategy will call this function, after every round of federated learning, passing two arguments: the current round of federated learning and the aggregated model parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def central_evaluate(server_round: int, arrays: ArrayRecord) -> MetricRecord:\n",
    "    \"\"\"Evaluate model on the server side.\"\"\"\n",
    "\n",
    "    # Load the model and initialize it with the received weights\n",
    "    model = Net(num_classes=10)\n",
    "    model.load_state_dict(arrays.to_torch_state_dict())\n",
    "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "    model.to(device)\n",
    "\n",
    "    # We use the testloader defined earlier in the centralized section\n",
    "    # Evaluate the model on the test set\n",
    "    loss, accuracy = test_fn(model, testloader, device)\n",
    "\n",
    "    # Return the evaluation metrics\n",
    "    return MetricRecord({\"accuracy\": accuracy, \"loss\": loss})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the above defined, we just need to wrap it all up in a `ServerApp` as we did earlier but this time using the `FedAvgCustom` that we just defined. We'll run the strategy for 10 rounds to observe the effect of LR decay.\n",
    "\n",
    "Remember we mentioned this `central_evaluate` will be called by the strategy. To do so we need to pass it to the strategy’s `start` method as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "9ad8dcea-8004-4c6e-a025-e168da636c88"
   },
   "outputs": [],
   "source": [
    "@server_app.main()\n",
    "def main(grid: Grid, context: Context) -> None:\n",
    "    \"\"\"Main entry point for the ServerApp.\"\"\"\n",
    "\n",
    "    # Load global model\n",
    "    global_model = Net(num_classes=10)\n",
    "    arrays = ArrayRecord(global_model.state_dict())\n",
    "\n",
    "    # Initialize FedAvgCustom strategy\n",
    "    strategy = FedAvgCustom(fraction_train=0.5, fraction_evaluate=1.0)\n",
    "\n",
    "    # Start strategy\n",
    "    result = strategy.start(\n",
    "        grid=grid,\n",
    "        initial_arrays=arrays,\n",
    "        train_config=ConfigRecord({\"lr\": 0.1}),\n",
    "        num_rounds=10,\n",
    "        evaluate_fn=central_evaluate,  # <--- here we pass the evaluate function\n",
    "    )\n",
    "\n",
    "    # Save final model to disk\n",
    "    print(\"\\nSaving final model to disk...\")\n",
    "    state_dict = result.arrays.to_torch_state_dict()\n",
    "    torch.save(state_dict, \"final_model.pt\")\n",
    "\n",
    "    # let's make a plot of the accuracy obtained after each round\n",
    "    # Note here we use the keys in central_evaluate\n",
    "    central_acc = [\n",
    "        rnd[\"accuracy\"] for rnd in result.evaluate_metrics_serverapp.values()\n",
    "    ]\n",
    "\n",
    "    plt.plot(central_acc, label=\"Accuracy\")\n",
    "    plt.xlabel(\"Round\")\n",
    "    plt.ylabel(\"Metric\")\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All that is left is to launch the simulation. You will note the following three differences compared to the previous run:\n",
    "\n",
    "- The global model is evaluated before round 1 and then after each round.\n",
    "- The learning rate is decreased by 50% at round 5\n",
    "- A line plot is generated at the end of the run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_simulation(\n",
    "    server_app=server_app, client_app=client_app, num_supernodes=NUM_PARTITIONS\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What's next?\n",
    "\n",
    "Well, if you enjoyed this content, consider giving us a ⭐️ on GitHub -> https://github.com/adap/flower\n",
    "\n",
    "* **[TUTORIAL]:** Check the `Flower Tutorial` for a more [step-by-step and comprehensive](https://flower.ai/docs/framework/tutorial-series-get-started-with-flower-pytorch.html) learning experience. How about running your Flower clients on the GPU? Find out how to do it in the [Flower Simulation Documentation](https://flower.ai/docs/framework/how-to-run-simulations.html)\n",
    "\n",
    "* Check the quickstart tutorials in https://flower.ai/docs/framework/, for example:\n",
    "    * [Quickstart PyTorch](https://flower.ai/docs/framework/tutorial-quickstart-pytorch.html)\n",
    "    * [Quickstart TensorFlow](https://flower.ai/docs/framework/tutorial-quickstart-tensorflow.html)\n",
    "    * [Quickstart JAX](https://flower.ai/docs/framework/tutorial-quickstart-jax.html)\n",
    "\n",
    "* Most examples in the Flower GitHub can run in simulation. These are some:\n",
    "\n",
    "    * [Advanced PyTorch](https://github.com/adap/flower/tree/main/examples/advanced-pytorch)\n",
    "    * [Finetuning a ViT](https://github.com/adap/flower/tree/main/examples/flowertune-vit)\n",
    "    * [Quickstart with 🤗 Huggingface](https://github.com/adap/flower/tree/main/examples/quickstart-huggingface)\n",
    "    * [Quickstart with XGBoost](https://github.com/adap/flower/tree/main/examples/xgboost-quickstart)\n",
    "    * [Quickstart with MLX](https://github.com/adap/flower/tree/main/examples/quickstart-mlx)\n",
    "\n",
    "\n",
    "* **[DOCS]** All Flower examples: https://flower.ai/docs/examples/\n",
    "\n",
    "* **[VIDEO]** Our Youtube channel: https://www.youtube.com/@flowerlabs\n",
    "\n",
    "Don't forget to join our Slack channel: https://flower.ai/join-slack/\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "3.11.10",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
